"""基于坑边缘拟合的陨石坑验证算法"""

import numpy as np
import cv2
from typing import Iterable
from utils.reader import Projection


def dem_crater_filter(
    nac_dem_prj_iter: Iterable[tuple[Projection, Projection]],
    k=1000,
    th=0.005,
    n=50,
):
    """
    根据DEM和NAC的对应位置，判断陨石坑位置，使用RANSAC算法完成拟合，拟合模型选取为椭圆抛物面
    a*x^2 + b*y^2 + c*x + d*y + e = z，其中中心坐标x0 = -c/2a, y0 = -d/2b
    两个轴向的半径长度为：rx = sqrt(1/(4a)), ry = sqrt(1/(4b))
    输入：
          n : 生成模型所需的最小数据集大小
          k : 最大迭代次数
          th : 残差容许阈值
          dem : 裁剪得到的DEM数据
          nac : 裁剪得到的NAC数据
    输出：
          (cnt_x, cnt_y) : 陨石坑中心坐标
          (rx, ry) : 陨石坑椭圆抛物面的两个轴向半径
          confidence : 当前位置内陨石坑的置信度
    """
    # 考虑用EM算法迭代，以获取最佳中心估计值和半径估计值.
    for nac_prj, dem_prj in nac_dem_prj_iter:
        dem = dem_prj.image
        nac = nac_prj.image
        # 首先生成坐标图，以便于拉直为一维数组
        x, y = np.meshgrid(range(dem.shape[0]), range(dem.shape[1]))
        x = x.flatten()
        y = y.flatten()
        z = dem.flatten() - dem.max()
        length = len(x)
        # 生成拟合数据集
        data = np.stack([x * x, y * y, x, y, np.ones_like(x)]).T
        # 获取初始采样值
        best_inliers = sample = np.random.choice(length, n, replace=False)
        best_num = 0
        for _ in range(k):
            # 生成模型
            model = np.linalg.lstsq(data[sample], z[sample], rcond=None)[0]
            # 计算残差
            inliers = np.abs(np.dot(data, model) - z) < th
            # 判断是否满足阈值
            if (inliers).sum() > n:
                sample = np.where(inliers)[0]
            else:
                sample = np.random.choice(length, n, replace=False)
            if (inliers).sum() >= best_num:
                best_num = (inliers).sum()
                best_inliers = inliers
        # 重新生成模型
        model = np.linalg.lstsq(data[best_inliers], z[best_inliers], rcond=None)[0]
        err = np.abs(np.dot(data, model) - z).sum()
        # 生成椭圆抛物面参数
        a, b, c, d, e = model
        # 计算中心坐标
        cntx = -c / (2 * a)
        cnty = -d / (2 * b)
        # 计算半径
        h = e - a * cntx**2 - b * cnty**2
        rx = np.sqrt(np.abs(h / (4 * a)))
        ry = np.sqrt(np.abs(h / (4 * b)))
        yield dem_prj.inverse(cntx, cnty), (rx, ry), err
